#!/usr/bin/python
# -*- coding:utf-8 -*-
# @FileName : DL6_test4_1.py
# Author    : myh

import torch
from d2l import torch as d2l

def corr2d_multi_in(X, K):
    # 先遍历“X”和“K”的第0个维度（通道维度），再把它们加在一起
    return sum(d2l.corr2d(x, k) for x, k in zip(X, K))

def corr2d_multi_in_out(X, K):
    # 迭代“K”的第0个维度，每次都对输入“X”执行互相关运算。
    # 最后将所有结果都叠加在一起
    return torch.stack([corr2d_multi_in(X, k) for k in K], 0)

def corr2d_multi_in_out_1x1(X, K):
    # 输入个数 长  宽
    c_i, h, w = X.shape
    # 输出个数
    c_o = K.shape[0]
    # 将输入的长宽放到一个维度中处理
    X = X.reshape((c_i, h * w))

    K = K.reshape((c_o, c_i))
    # 全连接层中的矩阵乘法
    Y = torch.matmul(K, X)
    return Y.reshape((c_o, h, w))


# X = torch.tensor([[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]],
#                [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]])
# K = torch.tensor([[[0.0, 1.0], [2.0, 3.0]], [[1.0, 2.0], [3.0, 4.0]]])
# # print(K.shape)
# # print(corr2d_multi_in(X, K))
#
# K = torch.stack((K, K + 1, K + 2), 0)
# print(K.shape)
# print(corr2d_multi_in_out(X, K))

# X = torch.normal(0, 1, (3, 3, 3))
# K = torch.normal(0, 1, (2, 3, 1, 1))
X = torch.tensor([[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]],
                  [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]],
                  [[2.0, 3.0, 4.0], [5.0, 6.0, 7.0], [8.0, 9.0, 10.0]]])
K = torch.tensor([[[1.0], [2.0], [3.0]],[[4.0],[5.0],[6.0]]])
#
Y1 = corr2d_multi_in_out_1x1(X, K)
print(Y1)
K = K.reshape((2,3,1,1))
print(K.shape)
Y2 = corr2d_multi_in_out(X, K)
print(Y2)
assert float(torch.abs(Y1 - Y2).sum()) < 1e-6